// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/sparksql/least_greatest.h>

#include <pollux/expression/eval_ctx.h>
#include <pollux/expression/expr.h>
#include <pollux/functions/sparksql/comparisons.h>
#include <pollux/type/type.h>

namespace kumo::pollux::functions::sparksql {
    namespace {
        template<typename Cmp, TypeKind kind>
        class LeastGreatestFunction final : public exec::VectorFunction {
            using T = typename TypeTraits<kind>::NativeType;

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                const size_t nargs = args.size();
                const auto nrows = rows.end();

                // Setup result vector.
                context.ensure_writable(rows, outputType, result);
                FlatVector<T> &flatResult = *result->as<FlatVector<T> >();

                // NULL all elements.
                rows.applyToSelected(
                    [&](vector_size_t row) { flatResult.set_null(row, true); });

                exec::LocalSelectivityVector cmpRows(context, nrows);
                exec::LocalDecodedVector decodedVectorHolder(context);
                // Column-wise process: one argument at a time.
                for (size_t i = 0; i < nargs; i++) {
                    decodedVectorHolder.get()->decode(*args[i], rows);

                    // Only compare with non-null elements of each argument
                    *cmpRows = rows;
                    if (auto *rawNulls = decodedVectorHolder->nulls(&rows)) {
                        cmpRows->deselectNulls(rawNulls, 0, nrows);
                    }

                    if (decodedVectorHolder->isIdentityMapping()) {
                        // Fast path: this argument is a FlatVector
                        cmpAndReplace(flatResult, *args[i]->as<FlatVector<T> >(), *cmpRows);
                    } else {
                        // Slow path: decode this argument
                        cmpAndReplace(flatResult, *decodedVectorHolder, *cmpRows);
                    }
                }
            }

            inline T getValue(const FlatVector<T> &v, vector_size_t i) const {
                return v.value_at(i);
            }

            inline T getValue(const DecodedVector &v, vector_size_t i) const {
                return v.value_at<T>(i);
            }

            template<typename VecType>
            void cmpAndReplace(
                FlatVector<T> &dst,
                const VecType &src,
                SelectivityVector &rows) const {
                const Cmp cmp;
                rows.applyToSelected([&](vector_size_t row) {
                    const auto srcVal = getValue(src, row);
                    if (dst.is_null_at(row) || cmp(srcVal, getValue(dst, row))) {
                        dst.set(row, srcVal);
                    }
                });
            }
        };

        template<template <typename> class Cmp>
        std::shared_ptr<exec::VectorFunction> makeImpl(
            const std::string &functionName,
            const std::vector<exec::VectorFunctionArg> &args) {
            POLLUX_CHECK_GE(args.size(), 2);
            for (size_t i = 1; i < args.size(); i++) {
                POLLUX_CHECK(*args[i].type == *args[0].type);
            }

            switch (args[0].type->kind()) {
#define SCALAR_CASE(kind)                            \
  case TypeKind::kind:                               \
    return std::make_shared<LeastGreatestFunction<   \
        Cmp<TypeTraits<TypeKind::kind>::NativeType>, \
        TypeKind::kind>>();
                SCALAR_CASE(BOOLEAN)
                SCALAR_CASE(TINYINT)
                SCALAR_CASE(SMALLINT)
                SCALAR_CASE(INTEGER)
                SCALAR_CASE(BIGINT)
                SCALAR_CASE(HUGEINT)
                SCALAR_CASE(REAL)
                SCALAR_CASE(DOUBLE)
                SCALAR_CASE(VARCHAR)
                SCALAR_CASE(VARBINARY)
                SCALAR_CASE(TIMESTAMP)
#undef SCALAR_CASE
                default:
                    POLLUX_NYI(
                        "{} does not support arguments of type {}",
                        functionName,
                        args[0].type->kind());
            }
        }
    } // namespace

    std::shared_ptr<exec::VectorFunction> makeLeast(
        const std::string &functionName,
        const std::vector<exec::VectorFunctionArg> &args,
        const core::QueryConfig & /*config*/) {
        return makeImpl<Less>(functionName, args);
    }

    std::vector<std::shared_ptr<exec::FunctionSignature> > leastSignatures() {
        std::vector<std::string> types = {
            "boolean",
            "tinyint",
            "smallint",
            "integer",
            "bigint",
            "real",
            "double",
            "varchar",
            "varbinary",
            "timestamp",
            "date"
        };
        std::vector<std::shared_ptr<exec::FunctionSignature> > signatures;

        for (const auto &type: types) {
            signatures.emplace_back(exec::FunctionSignatureBuilder()
                .returnType(type)
                .argumentType(type)
                .variableArity(type)
                .build());
        }
        signatures.emplace_back(exec::FunctionSignatureBuilder()
            .integerVariable("p")
            .integerVariable("s")
            .returnType("decimal(p,s)")
            .argumentType("decimal(p,s)")
            .variableArity("decimal(p,s)")
            .build());
        return signatures;
    }

    std::shared_ptr<exec::VectorFunction> makeGreatest(
        const std::string &functionName,
        const std::vector<exec::VectorFunctionArg> &args,
        const core::QueryConfig & /*config*/) {
        return makeImpl<Greater>(functionName, args);
    }

    std::vector<std::shared_ptr<exec::FunctionSignature> > greatestSignatures() {
        return leastSignatures();
    }
} // namespace kumo::pollux::functions::sparksql
